import torch
from torch import nn, std
from .utils import normal_init


class SampleProbNet(nn.Module):
    def __init__(self, hidden_dim=64, out_dim=64) -> None:
        super(SampleProbNet, self).__init__()

        act_func = nn.ReLU
        
        self.iou_net = nn.Sequential(
            nn.Linear(in_features=1, out_features=hidden_dim),
            act_func(),
            nn.Linear(in_features=hidden_dim, out_features=out_dim),
            act_func(),
        )

        self.prob_net = nn.Sequential(
            nn.Linear(in_features=80, out_features=hidden_dim * 2),
            act_func(),
            nn.Linear(in_features=hidden_dim*2, out_features=out_dim),
            act_func()
        )    

        self.cls_loss_net = nn.Sequential(
            nn.Linear(in_features=1, out_features=hidden_dim),
            act_func(),
            nn.Linear(in_features=hidden_dim, out_features=out_dim),
            act_func()
        )

        self.predictor = nn.Sequential(
            nn.Linear(in_features=out_dim * 3, out_features=out_dim),
            act_func(),
            nn.Linear(in_features=out_dim, out_features=1),
            nn.Sigmoid()
        )

        for m in [self.iou_net, self.prob_net, self.cls_loss_net]:
            normal_init(m[0], mean=0.0, std=0.0001, bias=0)
            normal_init(m[2], mean=0.0, std=0.0001, bias=0)

    def forward(self, ious, probs, cls_loss):
        iou_embed = self.iou_net(ious.view(-1, 1))
        prob_embed = self.prob_net(probs.view(-1, 80))
        cls_loss_embed = self.cls_loss_net(cls_loss.view(-1, 1))
        
        join_input = torch.cat((iou_embed, prob_embed, cls_loss_embed), dim=1)
        prob = self.predictor(join_input).view(-1)
        return prob
        
